package com.echo.AiManager;

import io.reactivex.Flowable;
import io.reactivex.Scheduler;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

public class StreamResponseService {

    /**
     * 通用流式返回方法，将数据逐步通过 SSE 发送给前端
     *
     * @param dataFlowable   数据流对象，包含逐步返回的数据
     * @param scheduler      线程调度器，可用于指定特定线程池（如 VIP 调度器）
     * @param dataProcessor  数据处理函数，将每个元素转换为所需的格式（如 JSON 字符串）
     * @param sseEmitter     SSE 连接对象，用于向前端推送数据
     * @throws Exception     处理过程中发生的异常
     */
    public <T> void sendStreamedResponse(Flowable<T> dataFlowable, 
                                         Scheduler scheduler, 
                                         Function<T, String> dataProcessor, 
                                         SseEmitter sseEmitter) throws Exception {
        
        // 计数器，用于判断是否形成完整数据块
        AtomicInteger counter = new AtomicInteger(0);
        
        // 用于拼接完整数据块
        StringBuilder stringBuilder = new StringBuilder();
        
        // 订阅数据流并处理逐个数据项
        dataFlowable
            .observeOn(scheduler)  // 指定线程调度器
            .map(dataProcessor::apply)  // 将数据转换为字符串（例如 JSON 格式）
            .map(message -> message.replaceAll("\\s", ""))  // 去除空白字符
            .filter(message -> !message.isEmpty())  // 过滤空消息
            .flatMap(message -> {
                List<Character> characterList = new ArrayList<>();
                for (char c : message.toCharArray()) {
                    characterList.add(c);
                }
                return Flowable.fromIterable(characterList);  // 按字符流式发送
            })
            .doOnNext(c -> {
                if (c == '{') counter.incrementAndGet();  // 增加计数器
                if (counter.get() > 0) stringBuilder.append(c);  // 拼接字符
                if (c == '}') {
                    counter.decrementAndGet();  // 减少计数器
                    if (counter.get() == 0) {
                        //todo： 当计数器归零时，表示数据块已完成，即发送给前端，这里是方法，上面是数据处理，可改动
                        sseEmitter.send(stringBuilder.toString());  // 发送数据到前端
                        stringBuilder.setLength(0);  // 清空缓冲区，准备下一次
                    }
                }
            })
            .doOnError(e -> sseEmitter.completeWithError(e))  // 错误处理
            .doOnComplete(sseEmitter::complete)  // 流完成时关闭 SSE
            .subscribe();
    }
}
